#include <stdio.h>
#include <stdlib.h>
#include <jpeglib.h>
#include <stdint.h>

// RGB to YUV conversion function
void rgb_to_yuv(uint8_t r, uint8_t g, uint8_t b, uint8_t *y, uint8_t *u, uint8_t *v) {
    *y = (uint8_t)(0.299 * r + 0.587 * g + 0.114 * b);
    *u = (uint8_t)(-0.14713 * r - 0.28886 * g + 0.436 * b + 128);
    *v = (uint8_t)(0.615 * r - 0.51499 * g - 0.10001 * b + 128);
}

// Function to convert JPG to YUV
int jpg_to_yuv(const char *jpg_filename, const char *yuv_filename, int width, int height) {
    struct jpeg_decompress_struct cinfo;
    struct jpeg_error_mgr jerr;
    FILE *infile, *outfile;
    JSAMPARRAY buffer;
    int row_stride;
    uint8_t *yuv_buffer;
    int y_size, uv_size;

    // Open JPG file
    if ((infile = fopen(jpg_filename, "rb")) == NULL) {
        fprintf(stderr, "Cannot open input JPG file %s\n", jpg_filename);
        return -1;
    }

    // Open YUV file
    if ((outfile = fopen(yuv_filename, "wb")) == NULL) {
        fprintf(stderr, "Cannot open output YUV file %s\n", yuv_filename);
        fclose(infile);
        return -1;
    }

    // Initialize JPEG decompression
    cinfo.err = jpeg_std_error(&jerr);
    jpeg_create_decompress(&cinfo);
    jpeg_stdio_src(&cinfo, infile);
    jpeg_read_header(&cinfo, TRUE);
    jpeg_start_decompress(&cinfo);

    // Check image dimensions
    width = cinfo.output_width;
    height = cinfo.output_height;
    printf("Image dimensions: %dx%d\n", width, height);

    // Calculate buffer sizes
    y_size = width * height;
    uv_size = (width / 2) * (height / 2);
    yuv_buffer = (uint8_t *)malloc(y_size + 2 * uv_size);

    // Allocate buffer for one row of pixel data
    row_stride = cinfo.output_width * cinfo.output_components;
    buffer = (*cinfo.mem->alloc_sarray)((j_common_ptr)&cinfo, JPOOL_IMAGE, row_stride, 1);

    // Process each row
    uint8_t *y_ptr = yuv_buffer;
    uint8_t *u_ptr = yuv_buffer + y_size;
    uint8_t *v_ptr = u_ptr + uv_size;
    int y_row = 0, uv_row = 0;

    while (cinfo.output_scanline < cinfo.output_height) {
        jpeg_read_scanlines(&cinfo, buffer, 1);

        // Convert RGB to YUV and downsample UV
        uint8_t *rgb_ptr = buffer[0];
        for (int x = 0; x < cinfo.output_width; x++) {
            uint8_t r = rgb_ptr[0];
            uint8_t g = rgb_ptr[1];
            uint8_t b = rgb_ptr[2];
            uint8_t y, u, v;

            rgb_to_yuv(r, g, b, &y, &u, &v);
            y_ptr[y_row * width + x] = y;

            // Downsample UV by 2x2
            if (y_row % 2 == 0 && x % 2 == 0) {
                u_ptr[uv_row * (width / 2) + (x / 2)] = u;
                v_ptr[uv_row * (width / 2) + (x / 2)] = v;
            }

            rgb_ptr += 3;
        }

        y_row++;
        if (y_row % 2 == 0) uv_row++;
    }

    // Write YUV data to file
    fwrite(y_ptr, 1, y_size, outfile);
    fwrite(u_ptr, 1, uv_size, outfile);
    fwrite(v_ptr, 1, uv_size, outfile);

    // Cleanup
    jpeg_finish_decompress(&cinfo);
    jpeg_destroy_decompress(&cinfo);
    fclose(infile);
    fclose(outfile);
    free(yuv_buffer);

    return 0;
}

int main(int argc, char *argv[]) {
    if (argc != 3) {
        fprintf(stderr, "Usage: %s <input.jpg> <output.yuv>\n", argv[0]);
        return -1;
    }

    int width = 0, height = 0;
    return jpg_to_yuv(argv[1], argv[2], width, height);
}